pacman::p_load(tidyverse, data.table, broom,modelsummary,kableExtra,latex2exp,
               ggpubr,lfe,fixest)
options("modelsummary_format_numeric_latex" = "plain")
rm(list=ls())
################################################################################

###### Data

spatial_id_a = 'county_id'
spatial_id_b = 'amc_id'

df = fread("inputs/data_supply_across.csv.gz")
df$y = df$farmland/df$forest
df = df[y>0,]
df$y = log(df$y)
df$x = log(df$payoff_nest)
df$iv = log(df$iv)
df_sample_xy = df[is.finite(x) & is.finite(y),]

df = df[is.finite(x) & is.finite(y) & is.finite(df$iv),]
df$location_fe = df$state_id
df$frontier=0
frontier_list = c('NORTE','NORDESTE','NOA','NEA','PATAGONIA','CUYO')
df[region %in% frontier_list,]$frontier=1

###### Regressions

m_ols = felm(y ~ x | year + location_fe | 0 | spatial_id, data=df)
m_ols_i = felm(y ~ x + x:frontier | year + location_fe | 0 | spatial_id, data=df)
m_iv = felm(y ~ 1 | year + location_fe | (x~iv) | spatial_id, data=df)
fs_iv = as.character(round(m_iv$stage1$iv1fstat$x[5],1))
m_iv_i = felm(y ~ 1 + x:frontier | year + location_fe | (x~iv) | spatial_id, data=df)
fs_iv_i = as.character(round(m_iv_i$stage1$iv1fstat$x[5],1))

###### Table

tab_title = 'Nested model - deforestation elasticity (substitution elasticity across nests).'
rows <- tribble(~term,               ~OLS, ~IV, ~OLS, ~IV,
                'Time FE'             ,'$\\checkmark$', '$\\checkmark$', '$\\checkmark$', '$\\checkmark$',
                'Location FE'         ,'$\\checkmark$', '$\\checkmark$', '$\\checkmark$', '$\\checkmark$',
                'KP-Wald First stage F-statistic'  ,'', fs_iv, '' , fs_iv_i)
attr(rows, 'position') <- c(5,6)
gm <- tibble::tribble(
  ~raw,        ~clean,          ~fmt,
  "nobs",      "Observations",  0)
final_table_tex <- msummary(list('OLS'=m_ols,'IV'=m_iv,'OLS'=m_ols_i,'IV'=m_iv_i),
                            title = paste0('\\label{tab: estimation - supply - across}',tab_title),
                            coef_omit = 'Intercept',
                            coef_map =  c('x'   = '$\\lambda$', 'x(fit)' ='$\\lambda$', 'x:frontier'='$\\lambda$ $\\times$ frontier region'),
                            estimate = "{estimate}{stars}",
                            gof_map = gm, 
                            escape=FALSE,
                            add_rows = rows,
                            output='latex')
final_table_tex = final_table_tex %>% kable_styling(latex_options = "HOLD_position")
kableExtra::save_kable(final_table_tex, file = "../../output/tables/table3b_across.tex")

###### Export parameters

spec_within = fread('inputs/spec_scale_within.csv')

lambda = m_iv_i$coefficients[2]
lambda_frontier =m_iv_i$coefficients[1]
lambda_est_core = round(lambda,4)
lambda_est_frontier =  round(lambda+lambda_frontier,4)

df = fread('../../data/landuse/clean/geographicunits_argentina.csv.gz')
setnames(df, old=spatial_id_a, new='spatial_id')
df_a = unique(df[, list(spatial_id, region)])
df = fread('../../data/landuse/clean/geographicunits_brazil.csv.gz')
setnames(df, old=spatial_id_b, new='spatial_id')
df_b = unique(df[, list(spatial_id, region)])
df = rbind(df_a,df_b)

df$lambda = lambda_est_core
df[region %in% frontier_list,]$lambda = lambda_est_frontier
df_across_elasticities = df
df_within_elasticities = fread("inputs/parameters_supply_within_tl.csv")

df=merge(df_across_elasticities, df_within_elasticities[,c('spatial_id','region','theta_lambda')], by=c('spatial_id','region'), all.x=T)
df$theta = df$theta_lambda*df$lambda
df$frontier = 0
df[region %in% frontier_list,]$frontier=1
df_params=df

rescale_lambda = unique(spec_within$rescale_lambda)
rescale_theta_lambda = unique(spec_within$rescale_theta_lambda)
df_params$theta_lambda = df_params$theta_lambda+rescale_theta_lambda
df_params$lambda = df_params$lambda+rescale_lambda
df_params$theta = df_params$theta_lambda*df_params$lambda
write.csv(df_params, "inputs/parameters_supply.csv", row.names = FALSE)

df=merge(df_sample_xy, df_params, by=c('spatial_id','region'), all.x=T)
df$log_AN = (1/df$theta_lambda)*df$x-(1/df$theta)*df$y
df$AN = exp(df$log_AN)
df = df[,c('year','country','spatial_id','region','state_id','frontier','log_AN','AN')]
write.csv(df, "inputs/parameters_supply_AN.csv", row.names = FALSE)

